import matplotlib.pyplot as plt
import torch
import torch.nn as nn

# 构造假数据，用于模拟测试
x = torch.linspace(0, 10, 100).unsqueeze(1)
y = torch.sin(x) * x
y += torch.normal(0, 0.4, y.shape)

# 1. 超参数
lr = 1e-2
Epochs = 10000


# 2. 创建模型
class Model(nn.Module):
    def __init__(self, input_sizes):
        super().__init__()
        self.fcs = nn.ModuleList()
        for i in range(len(input_sizes) - 1):
            if i != len(input_sizes) - 2:
                self.fcs.append(
                    nn.Linear(input_sizes[i], input_sizes[i + 1])
                )
                self.fcs.append(
                    nn.ReLU()
                )
            else:
                self.fcs.append(
                    nn.Linear(input_sizes[i], input_sizes[i + 1]),
                )

    def forward(self, x):
        for fc in self.fcs:
            x = fc(x)
        return x


input_sizes = [1, 128, 256, 128, 64, 1]
model = Model(input_sizes)

# 3. 损失函数
loss_fn = nn.MSELoss()

# 4. 优化器
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

'''准备画图'''
fig, (ax1, ax2) = plt.subplots(1, 2)

# 5. 循环训练
for epoch in range(Epochs):
    # 前向传播
    y_pre = model(x)
    # 计算损失
    loss = loss_fn(y_pre, y)
    # 反向传播
    loss.backward()
    # 更新参数
    optimizer.step()
    # 清空梯度
    optimizer.zero_grad()
    # 打印训练信息
    if epoch == 0 or (epoch + 1) % 100 == 0:
        print(f"[{epoch + 1}/{Epochs}] Loss:{loss.item():.4f}")
        ''' 更新绘图 '''
        ax1.clear()
        ax1.scatter(x, y, c="r")
        ax1.plot(x, y_pre.detach(), c="b")

        plt.pause(0.01)

# 6. 画图
plt.show()
